/**
   forest.cu - Created by Timothy Morey on 1/25/2012.

   This file contains the implementation for the cuda kernels in the forest fire
   simulator.
*/


#include <curand_kernel.h>

#include "forest.cuh"


////////////////////////////////////////////////////////////////////////////////
// CUDA Kernels:

__global__
void Downsample(unsigned int* forest, char* forestImg,
								unsigned short width, unsigned short height)
{
	int row = blockIdx.x;
	int col = threadIdx.x;

	while(col < width)
	{
		unsigned int cell = forest[row * width + col];
		char pix = 0;

		if(IsBurning(cell))
			pix |= 0xE0;

		char dgrowth = (char)((float)GetGrowth(cell) / 32.0f);
		pix |= 0x4 * dgrowth;
		
		forestImg[row * width + col] = pix;
		col += blockDim.x;
	}
}

__global__
void InitForest(unsigned int *grid, curandState* rngStates,
		unsigned short width, unsigned short height)
{
  int tid = blockIdx.x * blockDim.x + threadIdx.x;
  int row = blockIdx.x;
  int col = threadIdx.x;
  
  curandState rngState = rngStates[tid];

  while(col < width)
  {
    unsigned int val = 0xFF000000;
    val += RandomByte(rngState) * GROWTH_BUMP;
    grid[row * width + col] = val;
    col += blockDim.x;
  }

  rngStates[tid] = rngState;
}

__global__
void InitRNG(curandState* states, int seed)
{
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  curand_init(seed, i, 0, &states[i]);
}

__global__
void Step(unsigned int* curGrid, unsigned int* newGrid,
	  curandState* rngStates,
	  unsigned short width, unsigned short height)
{
  int tid = blockIdx.x * blockDim.x + threadIdx.x;
  int row = blockIdx.x;
  int col = threadIdx.x;

  curandState rngState = rngStates[tid];

  while(col < width)
  {
    unsigned int localRegion[9];
    FillLocalRegion(curGrid, width, height, row, col, localRegion);
    
    unsigned char burningCount = 0;
    for(int i = 0; i < 9; i++)
      if(IsBurning(localRegion[i]))
				burningCount++;
    
    unsigned int newVal = localRegion[4];
    
    if(IsBurning(localRegion[4]))
    {
      if(GetGrowth(localRegion[4]) > 0)
				Burn(newVal);
      else
				Extinguish(newVal);
    }
    else
    {
      if(IsGround(localRegion[4]))
      {
				if(FlipCoin(GROUND_GROWTH_PROB, rngState))
					Grow(newVal, rngState);
      }
      else if(burningCount > 0 &&
							FlipCoin(burningCount * FIRE_SPREAD_PROB, rngState))
      {
				Ignite(newVal);
      }
      else if(burningCount == 0 &&
							FlipCoin(RANDOM_SPARK_PROB, rngState))
      {
				Ignite(newVal);
      }
      else if(IsGrass(localRegion[4]))
      {
				Grow(newVal, rngState);
      }
      else if(IsTrees(localRegion[4]))
      {
				Grow(newVal, rngState);
      }
    }
		
    newGrid[row*width+col] = newVal;
    col += blockDim.x;
  }
	
  rngStates[tid] = rngState;
}


////////////////////////////////////////////////////////////////////////////////
// Device Helper Functions:

__device__
__host__
void Burn(unsigned int& cellValue)
{
  if(GetGrowth(cellValue) > 0)
    cellValue -= GROWTH_BUMP;
}

__device__
__host__
void FillLocalRegion(unsigned int* grid, unsigned short width, unsigned short height,
		     unsigned short row, unsigned short col, unsigned int* localRegion)
{
  for(int i = 0; i < 9; i++)
    localRegion[i] = 0;

  if(row > 0 && col > 0)
    localRegion[0] = grid[GetForestIndex(row-1, width, col-1)];

  if(row > 0)
    localRegion[1] = grid[GetForestIndex(row-1, width, col)];

  if(row > 0 && col + 1 < width)
    localRegion[2] = grid[GetForestIndex(row-1, width, col+1)];

  if(col > 0)
    localRegion[3] = grid[GetForestIndex(row, width, col-1)];
  
  localRegion[4] = grid[GetForestIndex(row, width, col)];

  if(col + 1 < width)
    localRegion[5] = grid[GetForestIndex(row, width, col+1)];

  if(row + 1 < height && col > 0)
    localRegion[6] = grid[GetForestIndex(row+1, width, col-1)];

  if(row + 1 < height)
    localRegion[7] = grid[GetForestIndex(row+1, width, col)];

  if(row + 1 < height && col + 1 < width)
    localRegion[8] = grid[GetForestIndex(row+1, width, col+1)];
}

__device__
bool FlipCoin(float probability, curandState& rngState)
{
  return RandomFloat(rngState) < probability;
}

__device__
__host__
int GetForestIndex(int row, int width, int col)
{
  return row * width + col;
}

__device__
__host__
unsigned char GetGrowth(unsigned int cellValue)
{
  return (unsigned char)((cellValue & GROWTH_CHANNEL) >> 8);
}

__device__
void Grow(unsigned int& cellValue, curandState& rngState)
{
  unsigned char growth = GetGrowth(cellValue);

  if(growth < GRASS_MAX && 
     FlipCoin(GRASS_GROWTH_PROB, rngState))
  {
    cellValue += GROWTH_BUMP;
  }
  else if(growth < TREE_MAX &&
	  FlipCoin(TREE_GROWTH_PROB, rngState))
  {
    cellValue += GROWTH_BUMP;
  }
}

__device__
__host__
void Ignite(unsigned int& cell)
{
  cell |= FIRE;
}

__device__
__host__
void Extinguish(unsigned int& cell)
{
  cell &= 0xFFFFFF00;
}

__device__
__host__
bool IsBurning(unsigned int cell)
{
  return cell & FIRE;
}

__device__
__host__
bool IsGround(unsigned int cellValue)
{
  return GetGrowth(cellValue) == 0;
}

__device__
__host__
bool IsGrass(unsigned int cellValue)
{
  unsigned char growth = GetGrowth(cellValue);
  return growth >= GRASS_MIN && growth <= GRASS_MAX;
}

__device__
__host__
bool IsTrees(unsigned int cellValue)
{
  unsigned char growth = GetGrowth(cellValue);
  return growth >= TREE_MIN;
}

__device__
float RandomFloat(curandState& rngState)
{
  return curand_uniform(&rngState);
}

__device__
unsigned int RandomInt(curandState& rngState)
{
  return curand(&rngState);
}

__device__
unsigned char RandomByte(curandState& rngState)
{
  return (unsigned char)(RandomInt(rngState) % 256);
}
